package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// LogMessage 定义日志消息结构体
type LogMessage struct {
	B  string `json:"b"`  // 设备名称
	T  string `json:"t"`  // 类型
	M  string `json:"m"`  // 日志内容
	Ti string `json:"ti"` // 时间
}

var logFiles = make(map[string]*os.File)
var mu sync.Mutex

// 端口配置，默认 8080
var port = flag.String("port", "8080", "服务器监听端口")

// 是否写入日志，默认开启
var writeLog = flag.Bool("write-log", false, "是否写入日志文件")

func main() {
	flag.Parse()

	// 确保日志目录存在
	if *writeLog {
		if err := os.MkdirAll("logs", 0755); err != nil {
			log.Fatal("创建日志目录失败:", err)
		}
	}

	http.HandleFunc("/ws", handleWebSocket)
	http.HandleFunc("/", serveHome)
	http.Handle("/logs/", http.StripPrefix("/logs/", http.FileServer(http.Dir("logs"))))

	log.Println("服务器启动在 :" + *port)
	log.Fatal(http.ListenAndServe(":"+*port, nil))
}

// serveHome 处理首页请求
func serveHome(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, `
	<!DOCTYPE html>
	<html>
	<head>
		<title>日志查看器</title>
		<style>
			body { font-family: Arial, sans-serif; margin: 20px; }
			#logs { 
				height: 80vh; 
				overflow-y: auto; 
				border: 1px solid #ddd; 
				padding: 10px; 
				background: #f9f9f9;
			}
			#logs div { 
				margin: 5px 0; 
				padding: 5px; 
				border-bottom: 1px solid #eee; 
			}
			#status { 
				padding: 5px; 
				background: #eee; 
				margin-bottom: 10px;
			}
		</style>
	</head>
	<body>
		<h1>WebSocket 日志服务</h1>
		<div id="status">连接状态: <span id="statusText">正在连接...</span></div>
		<div id="logs"></div>
		<script>
			const logsDiv = document.getElementById('logs');
			const statusText = document.getElementById('statusText');
			const ws = new WebSocket('ws://'+location.host+'/ws');
			
			// 标记是否应该自动滚动
			let shouldAutoScroll = true;
			
			// 监听滚动事件
			logsDiv.addEventListener('scroll', function() {
				const threshold = 10; // 容错阈值
				shouldAutoScroll = logsDiv.scrollHeight - logsDiv.scrollTop - logsDiv.clientHeight <= threshold;
			});
			
			ws.onopen = function() {
				statusText.textContent = '已连接';
				statusText.style.color = 'green';
			};
			
			ws.onclose = function() {
				statusText.textContent = '已断开';
				statusText.style.color = 'red';
			};
			
			ws.onmessage = function(e) {
				const log = JSON.parse(e.data);
				const logDiv = document.createElement('div');
				logDiv.innerHTML = '<strong>[' + log.ti + '] ' + log.b + ' ' + log.t + '</strong>: ' + log.m;
				logsDiv.appendChild(logDiv);
				
				// 只有当前在底部时才自动滚动
				if(shouldAutoScroll) {
					logsDiv.scrollTop = logsDiv.scrollHeight;
				}
			};
		</script>
	</body>
	</html>
	`)
}

// 添加全局变量来存储所有客户端连接
var clients = make(map[*websocket.Conn]bool)
var clientsMu sync.Mutex

// handleWebSocket 处理 WebSocket 连接
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket升级失败:", err)
		return
	}
	defer conn.Close()

	// 注册新客户端
	clientsMu.Lock()
	clients[conn] = true
	clientsMu.Unlock()

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println("读取消息失败:", err)
			break
		}

		println(string(msg))
		// 解析日志消息
		var logMsg LogMessage
		if err := json.Unmarshal(msg, &logMsg); err != nil {
			log.Println("解析日志消息失败:", err)
			continue
		}
		logMsg.Ti = time.Now().Format("2006-01-02 15:04:05")
		// 写入日志文件
		if *writeLog {
			if err := writeLogToFile(logMsg); err != nil {
				log.Println("写入日志文件失败:", err)
			}
		}

		// 广播消息给所有客户端
		clientsMu.Lock()
		for client := range clients {
			if err := client.WriteJSON(logMsg); err != nil {
				log.Println("广播消息失败:", err)
				client.Close()
				delete(clients, client)
			}
		}
		clientsMu.Unlock()
	}
}

// writeLogToFile 将日志消息写入文件
func writeLogToFile(logMsg LogMessage) error {
	mu.Lock()
	defer mu.Unlock()

	// 获取或创建日志文件
	file, ok := logFiles[logMsg.B]
	if !ok {
		logPath := filepath.Join("logs", logMsg.B+".log")
		f, err := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			return err
		}
		logFiles[logMsg.B] = f
		file = f
	}

	// 写入日志
	_, err := file.WriteString(fmt.Sprintf("[%s] %s %s: %s\n",
		logMsg.Ti,
		logMsg.B,
		logMsg.T,
		logMsg.M))

	return err
}
